#!/usr/bin/env python3
import rospy
import cv2
import os
import glob
from cv_bridge import CvBridge
from sensor_msgs.msg import Image

class ImageSaver:
    def __init__(self):
        self.bridge = CvBridge()

        # 图像保存目录
        self.color_dir = rospy.get_param('~color_save_path', '/root/ros_ws/src/OrbbecSDK_ROS1/scripts/data/Color2')
        self.depth_dir = rospy.get_param('~depth_save_path', '/root/ros_ws/src/OrbbecSDK_ROS1/scripts/data/Depth2')
        os.makedirs(self.color_dir, exist_ok=True)
        os.makedirs(self.depth_dir, exist_ok=True)

        # 订阅话题
        rospy.Subscriber('/camera/color/image_raw', Image, self.color_callback)
        rospy.Subscriber('/camera/depth/image_raw', Image, self.depth_callback)

        self.color_count = 0
        self.depth_count = 0

        # 设置定时器，每60秒清理一次图片
        self.cleanup_timer = rospy.Timer(rospy.Duration(60.0), self.cleanup_images)

        rospy.loginfo("ImageSaver node started, saving to:")
        rospy.loginfo("Color images: %s", self.color_dir)
        rospy.loginfo("Depth images: %s", self.depth_dir)
        rospy.loginfo("Images will be cleaned up every 60 seconds")

    def cleanup_images(self, event):
        """每60秒清理一次存储的图片"""
        try:
            # 清理彩色图片
            color_files = glob.glob(os.path.join(self.color_dir, "*.png"))
            for file in color_files:
                os.remove(file)
            
            # 清理深度图片
            depth_files = glob.glob(os.path.join(self.depth_dir, "*.png"))
            for file in depth_files:
                os.remove(file)
            
            # 重置计数器
            self.color_count = 0
            self.depth_count = 0
            
            rospy.loginfo(f"Cleaned up {len(color_files)} color images and {len(depth_files)} depth images")
        except Exception as e:
            rospy.logerr(f"Failed to cleanup images: {e}")

    def color_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            filename = os.path.join(self.color_dir, f"color_{self.color_count:04d}.png")
            cv2.imwrite(filename, cv_image)
            rospy.loginfo(f"Saved color image: {filename}")
            self.color_count += 1
        except Exception as e:
            rospy.logerr(f"Failed to save color image: {e}")

    def depth_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
            filename = os.path.join(self.depth_dir, f"depth_{self.depth_count:04d}.png")
            # OpenCV 只支持 16-bit PNG 深度图保存
            cv2.imwrite(filename, cv_image)
            rospy.loginfo(f"Saved depth image: {filename}")
            self.depth_count += 1
        except Exception as e:
            rospy.logerr(f"Failed to save depth image: {e}")

if __name__ == '__main__':
    rospy.init_node('image_saver', anonymous=True)
    saver = ImageSaver()
    rospy.spin()
